package health_system.health_system.config;

import health_system.health_system.security.JwtTokenProvider;
import jakarta.websocket.HandshakeResponse;
import jakarta.websocket.server.HandshakeRequest;
import jakarta.websocket.server.ServerEndpointConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.security.core.Authentication;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.util.List;
import java.util.Map;

/**
 * WebSocket配置器
 * 处理WebSocket握手过程，进行认证和初始化
 * 整合了JWT认证和用户ID提取功能
 */
@Component
public class WebSocketConfigurator extends ServerEndpointConfig.Configurator implements ApplicationContextAware {
    private static final Logger logger = LoggerFactory.getLogger(WebSocketConfigurator.class);
    private static volatile BeanFactory context;

    @Override
    public <T> T getEndpointInstance(Class<T> clazz) throws InstantiationException {
        return context.getBean(clazz);
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        WebSocketConfigurator.context = applicationContext;
    }
    
    @Override
    public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
        try {
            // 添加CORS头以允许WebSocket连接
            response.getHeaders().put("Access-Control-Allow-Origin", List.of("*"));
            response.getHeaders().put("Access-Control-Allow-Methods", List.of("GET, POST, OPTIONS"));
            response.getHeaders().put("Access-Control-Allow-Headers", List.of("Content-Type, Authorization"));
            response.getHeaders().put("Access-Control-Allow-Credentials", List.of("true"));
            response.getHeaders().put("Access-Control-Max-Age", List.of("3600"));
            
            logger.debug("WebSocket handshake: Added CORS headers");
            
            Map<String, List<String>> headers = request.getHeaders();
            
            // 1. 从请求中提取用户ID
            String userId = extractUserId(request);
            
            if (StringUtils.hasText(userId)) {
                // 将用户ID存储在会话属性中
                sec.getUserProperties().put("userId", userId);
                logger.debug("WebSocket handshake: User ID set to {}", userId);
            } else {
                logger.debug("WebSocket handshake: No user ID found");
            }
            
            // 2. 提取JWT令牌并进行验证
            if (headers.containsKey("Authorization")) {
                List<String> authHeaders = headers.get("Authorization");
                if (authHeaders != null && !authHeaders.isEmpty()) {
                    String authHeader = authHeaders.get(0);
                    if (StringUtils.hasText(authHeader) && authHeader.startsWith("Bearer ")) {
                        String token = authHeader.substring(7);
                        sec.getUserProperties().put("token", token);
                        
                        try {
                            // 从上下文获取JWT提供程序
                            if (context != null) {
                                JwtTokenProvider tokenProvider = context.getBean(JwtTokenProvider.class);
                                if (tokenProvider != null && tokenProvider.validateToken(token)) {
                                    Authentication auth = tokenProvider.getAuthentication(token);
                                    sec.getUserProperties().put("authentication", auth);
                                    logger.info("JWT authentication successful for WebSocket handshake");
                                } else {
                                    logger.warn("Invalid JWT token in WebSocket handshake");
                                }
                            }
                        } catch (Exception e) {
                            logger.error("Error validating JWT token in WebSocket handshake: {}", e.getMessage());
                        }
                    }
                }
            }
            
            // 3. 设置允许的子协议（如果需要）
            if (headers != null && headers.containsKey("Sec-WebSocket-Protocol")) {
                List<String> requestedProtocols = headers.get("Sec-WebSocket-Protocol");
                if (requestedProtocols != null && !requestedProtocols.isEmpty()) {
                    String selectedProtocol = requestedProtocols.get(0); // 简单起见，选择第一个
                    response.getHeaders().put("Sec-WebSocket-Protocol", List.of(selectedProtocol));
                    logger.debug("WebSocket handshake: Selected subprotocol {}", selectedProtocol);
                }
            }
        } catch (Exception e) {
            logger.error("WebSocket handshake modification failed: {}", e.getMessage(), e);
        }
        
        // 调用父类方法完成握手
        super.modifyHandshake(sec, request, response);
    }
    
    /**
     * 从请求中提取用户ID
     * @param request 握手请求
     * @return 用户ID，如果不存在则返回null
     */
    private String extractUserId(HandshakeRequest request) {
        // 从查询参数中提取
        Map<String, List<String>> parameters = request.getParameterMap();
        if (parameters != null && parameters.containsKey("userId")) {
            List<String> values = parameters.get("userId");
            if (values != null && !values.isEmpty()) {
                return values.get(0);
            }
        }
        
        // 如果查询参数中没有，尝试从URI中提取
        String requestUri = null;
        if (request.getRequestURI() != null) {
            requestUri = request.getRequestURI().toString();
        }
        
        if (requestUri != null && requestUri.contains("userId=")) {
            String[] parts = requestUri.split("userId=");
            if (parts.length >= 2) {
                String userId = parts[1];
                int ampIndex = userId.indexOf('&');
                if (ampIndex != -1) {
                    userId = userId.substring(0, ampIndex);
                }
                return userId;
            }
        }
        
        return null;
    }
    
    @Override
    public boolean checkOrigin(String originHeaderValue) {
        // 这里可以实现对WebSocket连接源的检查
        // 简单起见，我们允许所有源
        return true;
    }
} 